package com.pearl.sharding.level.table.demo.algorithm;

import lombok.Getter;
import org.apache.shardingsphere.infra.binder.statement.SQLStatementContext;
import org.apache.shardingsphere.infra.binder.statement.dml.DeleteStatementContext;
import org.apache.shardingsphere.infra.binder.statement.dml.UpdateStatementContext;
import org.apache.shardingsphere.infra.executor.audit.exception.SQLAuditException;
import org.apache.shardingsphere.infra.metadata.database.ShardingSphereDatabase;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.infra.metadata.user.Grantee;
import org.apache.shardingsphere.sharding.spi.ShardingAuditAlgorithm;
import org.apache.shardingsphere.sql.parser.sql.dialect.statement.mysql.dml.MySQLDeleteStatement;
import org.apache.shardingsphere.sql.parser.sql.dialect.statement.mysql.dml.MySQLUpdateStatement;

import java.util.List;
import java.util.Properties;

/**
 * @author TD
 * @version 1.0
 * @date 2024/2/29
 */
public final class LimitRequiredShardingAuditAlgorithm implements ShardingAuditAlgorithm {

    @Getter
    private Properties props;
    @Override
    public void init(final Properties props) {
        this.props = props;
    }

    @Override
    public void check(SQLStatementContext sqlStatementContext, List<Object> list, Grantee grantee, ShardingSphereRuleMetaData shardingSphereRuleMetaData, ShardingSphereDatabase shardingSphereDatabase) {
        if (sqlStatementContext instanceof UpdateStatementContext && ((MySQLUpdateStatement) sqlStatementContext.getSqlStatement()).getLimit().isEmpty()) {
            throw new SQLAuditException("Not allow update without limit");
        }
        if (sqlStatementContext instanceof DeleteStatementContext && ((MySQLDeleteStatement) sqlStatementContext.getSqlStatement()).getLimit().isEmpty()) {
            throw new SQLAuditException("Not allow delete without limit");
        }
    }

    @Override
    public String getType() {
        return "LIMIT_REQUIRED";
    }
}
